package assignment01;

import java.util.ArrayList;
import java.util.Collection;
import java.util.Comparator;
import java.util.HashSet;
import java.util.Iterator;
import java.util.List;
import java.util.Set;
import java.util.SortedSet;
import java.util.TreeSet;

import nlp.langmodel.LanguageModel;
import nlp.util.Counter;

public class UnigramModel implements LanguageModel {
	private Collection<List<String>> trainingCollection;
	private Collection<List<String>> validationCollection;
	private Counter<String> wordCounter = new Counter<String>();
	private Counter<Integer> nCounter = new Counter<Integer>();
	private static final double INCREMENT_FACTOR = 1.0;
	private final static Integer THRESHOLD = 5;
	private Double[] N = new Double[THRESHOLD + 2];
	private Double[] C = new Double[THRESHOLD + 1];
	private static final String STOP = "</S>";
	private double vacabulary = 0.0;
	private double unseenCount = 0.0;
	private double seenCount = 0.0;

	public UnigramModel(Collection<List<String>> trainingCollection,
			Collection<List<String>> validationCollection) {
		this.trainingCollection = trainingCollection;
		this.validationCollection = validationCollection;
		train();
		this.vacabulary = wordCounter.keySet().size();
		this.unseenCount = 1.0;
		this.seenCount = wordCounter.totalCount();
		System.out.println("before smoothing, total token is " + seenCount);
		smooth();
		System.out.println("C[0] is " + C[0]);
		System.out.println("after smoothing, total token is " + validate() + ", and N[1] is " + N[1]);
		linearRegression();
	}

	public double getVacabulary() {
		return vacabulary;
	}

	public void setVacabulary(double vacabulary) {
		this.vacabulary = vacabulary;
	}

	void train() {
		for (List<String> sentence : trainingCollection) {
			List<String> stopSentence = new ArrayList<String>(sentence);
			stopSentence.add(STOP);

			for (String word : stopSentence) {
				wordCounter.incrementCount(word, INCREMENT_FACTOR);
			}
		}
		
		// find count of counts
		for (String word : wordCounter.keySet()) {
			double count = wordCounter.getCount(word);
			nCounter.incrementCount((int)count, INCREMENT_FACTOR);
		}
		
	}
	
	void smooth() {
		Set<Integer> treeSet = new TreeSet<Integer>();
		treeSet.addAll(nCounter.keySet());
		Iterator<Integer> iter = treeSet.iterator();
		
		int next = 0;
		while (iter.hasNext() && next <= THRESHOLD) {
			next = iter.next();
			N[next] = nCounter.getCount(next);
		}
		N[0] = unseenCount;
		
		
		for(int c = 1; c <= THRESHOLD; c++) {
			double cStar = Math.log((c + 1) * N[c + 1] / N[c]
					- (c * (THRESHOLD + 1) * N[THRESHOLD + 1]) / N[1])
					- Math.log(1.0 - (THRESHOLD + 1) * N[THRESHOLD + 1] / N[1]);
			cStar = Math.pow(Math.E, cStar);
			C[c] = cStar;
		}
		C[0] = N[1] / unseenCount;
	}
	
	double validate() {
		double sum = 0.0;
		for (String word : wordCounter.keySet()) {
			double count = wordCounter.getCount(word);
			if ((int)count <= THRESHOLD) {
				sum += C[(int)count];
			} else {
				sum += count;
			}
		}
		return sum;
	}
	
	
	
	
	
	
	
	void linearRegression() {
		
		
		Set<Double> valueSet = new HashSet<Double>();
		for (Integer key : nCounter.keySet()) {
			valueSet.add(nCounter.getCount(key));
		}
		
		double aveX = calIntegerMean(nCounter.keySet());
		double aveY = calDoubleMean(valueSet);
		double b = calPowerLaw(aveX, aveY, nCounter);
		double log_a = aveY - b * aveX;
		System.out.println("the size of N is " + nCounter.size());
		System.out.println("Linear regression results: log_a = " + log_a + "; b = " + b);
	}

	static double calIntegerMean(Set<Integer> set) {
		double sum = 0;
		for (Integer key : set) {
			sum += key;
		}
//		System.out.println(sum);
		return sum / set.size();
	}
	
	static double calDoubleMean(Set<Double> set) {
		double sum = 0;
		for (Double key : set) {
			sum += key;
		}
//		System.out.println(sum);
		return sum / set.size();
	}
	
	static double calPowerLaw(double aveX, double aveY, Counter<Integer> nCounter) {
		double numerator = 0.0;
		double denominator = 0.0;
		for (Integer x : nCounter.keySet()) {
			numerator += (x - aveX) * (nCounter.getCount(x) - aveY);
			denominator += Math.pow(x - aveX, 2); 
		}
		return numerator / denominator;
	}
	
	static void averagingTransform(Counter<Integer> nCounter) {
		Set<Integer> treeSet = new TreeSet<Integer>();
		treeSet.addAll(nCounter.keySet());
		Iterator<Integer> iter = treeSet.iterator();
		Integer previous = 0;
		Integer current = iter.next();
		Integer next = 0;
		while (iter.hasNext()) {
			next = iter.next();
			System.out.println("treeset key " + current);
			Double n = nCounter.getCount(current);
			if (previous >= next) {
				System.out.println("shit");
			} else if (previous != 0){
				nCounter.setCount(current, 2 * n / (next - previous));
				System.out.println("fuck, " + "previous is " + previous + ", next is " + next);
			} else {
				System.out.println("shit, " + "previous is " + previous + ", next is " + next);
			}
			previous = current;
			current = next;
		}
	}
	
	double getWordLogProbability(List<String> sentence, int index) {
		String word = sentence.get(index);
		return getPWord(word);
	}

	/**
	 * @param word
	 * @return
	 */
	double getPWord(String word) {
		double count = 0.0;
		int intCount = (int) wordCounter.getCount(word);
		if (intCount <= THRESHOLD) {
			count = C[intCount];
		} else {
			count = intCount;
		}

//		System.out.println("word probability is " + Math.log(count / unigramTotal) / Math.log(2.0));
		return Math.log(count / seenCount) / Math.log(2.0);
	}
	
	@Override
	public double getSentenceProbability(List<String> sentence) {
		List<String> stoppedSentence = new ArrayList<String>(sentence);
		stoppedSentence.add(STOP);
		double probability = 0.0;
		for (int i = 0; i < stoppedSentence.size(); i++) {
//			System.out.println("the log word probability is " + Math.log(getWordProbability(stoppedSentence, i)));
			probability += getWordLogProbability(stoppedSentence, i);
		}
//		return probability;
		return Math.pow(2.0, probability);
	}
	
	double getWordCount(String word) {
		return wordCounter.getCount(word);
	}
	String generateWord() {
		double sample = Math.random();
	    double sum = 0.0;
	    for (String word : wordCounter.keySet()) {
	    	double count = 0.0;
			int intCount = (int) wordCounter.getCount(word);
			if (intCount <= THRESHOLD) {
				count = C[intCount];
			} else {
				count = intCount;
			}	
	      sum += count / seenCount;
	      if (sum > sample) {
	        return word;
	      }
	    }
	    return "*UNKNOWN*";
	}
	
	@Override
	public List<String> generateSentence() {
		List<String> sentence = new ArrayList<String>();
	    String word = generateWord();
	    while (!word.equals(STOP)) {
	      sentence.add(word);
	      word = generateWord();
	    }
	    return sentence;
	}
	
	public Counter<String> getUnigramCounter() {
		return this.wordCounter;
	}
}
